import networkx as nx

from algorithm.simplified_graph.compress_graph import compress_graph


def PRTH(G):
    # 重新计算PageRank值
    pr = nx.pagerank(G)
    # 计算边权重并保存到边属性中
    for u, v in G.edges():
        weight = pr[u] / (pr[u] + pr[v])
        G[u][v]['weight'] = max(weight, 1 - weight)
    # 计算阈值
    edge_weights = nx.get_edge_attributes(G, 'weight')
    th = sum(edge_weights.values()) / len(G.edges())
    compress_graph(G)
    # 过滤小于阈值的边
    for u, v in G.edges():
        if G[u][v]['weight'] <= th:
            G.remove_edge(u, v)
    # 重新计算PageRank值
    pr = nx.pagerank(G)
    # 对PageRank值进行排序，得到种子节点
    seed_nodes = sorted(pr, key=lambda x: pr[x], reverse=True)
    return seed_nodes
